import java.io.*;
import java.util.*;


public class HuffmanCode {
    public static void main(String[] args) {
        //String content = "i like like like java do you like a java";
        //byte[] bytes = content.getBytes();
        //System.out.println(Arrays.toString(bytes));
        //byte[] bytes1 = huffmanZip(bytes);
        //System.out.println(Arrays.toString(bytes1));
        //byte[] decode = decode(huffmanCodes, bytes1);
        //System.out.println(new String(decode));
        String srcFile = "d://a.bmp";
        String dstFile = "d://dst.xlzip";
        zipFile(srcFile, dstFile);
        String unzipFile = "d://dst.bmp";
        unZipFile(dstFile, unzipFile);
    }

    /**
     * 压缩文件
     *
     * @param srcFile 你希望压缩的文件的全路径
     * @param dstFile 压缩后的目录
     */
    public static void zipFile(String srcFile, String dstFile) {
        FileInputStream is = null;
        OutputStream os = null;
        //对象输出流
        ObjectOutputStream oos = null;
        try {
            //创建文件输入流
            is = new FileInputStream(srcFile);
            //创建一个和源文件大小一样的byte数组
            byte[] bytes = new byte[is.available()];
            //读取is流
            is.read(bytes);
            //对源文件进行压缩
            byte[] huffmanBytes = huffmanZip(bytes);
            //创建文件的输出流
            os = new FileOutputStream(dstFile);
            //创建一个和文件输出流关联的ObjectOutPutStream
            oos = new ObjectOutputStream(os);
            //以对象流的方式写入赫夫曼编码,为了以后恢复文件使用
            //把huffmanBytes数组写入到对象输出流
            oos.writeObject(huffmanBytes);
            //把赫夫曼编码表也写入
            oos.writeObject(huffmanCodes);


        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            try {
                oos.close();
                os.close();
                is.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }


    //解压文件出现过一次解压bytes越界的情况,不知道啥原因

    /**
     * 解压文件
     *
     * @param zipFile 要解压的文件
     * @param dstFile 解压到的路径
     */
    public static void unZipFile(String zipFile, String dstFile) {
        InputStream is = null;
        ObjectInputStream ois = null;
        OutputStream os = null;
        try {
            is = new FileInputStream(zipFile);
            ois = new ObjectInputStream(is);
            //读取byte数组
            byte[] huffmanBytes = (byte[]) ois.readObject();
            //读取编码表
            Map<Byte, String> huffmanCodes = (Map<Byte, String>) ois.readObject();

            //解码
            byte[] bytes = decode(huffmanCodes, huffmanBytes);
            //创建文件输出流
            os = new FileOutputStream(dstFile);
            os.write(bytes);
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            try {
                os.close();
                ois.close();
                is.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }


    }


    //封装整个方法

    /**
     * @param bytes 原始字节数组
     * @return 压缩后的数组
     */
    public static byte[] huffmanZip(byte[] bytes) {
        //将bytes转换成nodes，并生成一个赫夫曼树
        Node root = createHuffManTtree(getNodes(bytes));
        //将赫夫曼树生成对应的编码表
        Map<Byte, String> codes = getCodes(root);
        // 压缩原始bytes数组
        byte[] huffmanBytes = zip(bytes, codes);
        return huffmanBytes;
    }


    /**
     * @param huffmanCodes 哈夫曼编码表
     * @param huffmanBytes 哈法曼编码的数组
     * @return 解压缩, 原字节数组
     */
    public static byte[] decode(Map<Byte, String> huffmanCodes, byte[] huffmanBytes) {
        StringBuilder stringBuilder = new StringBuilder();
        int index = 0;
        for (; index < huffmanBytes.length - 1; index++) {
            stringBuilder.append(byteToBitString(true, huffmanBytes[index]));
        }
        //获得最后一个字节
        byte a = huffmanBytes[index];
        //如果最后一个字节是正数就不用向高位补8个0
        if (a > 0) {
            stringBuilder.append(byteToBitString(false, a));
        } else {
            stringBuilder.append(byteToBitString(true, a));
        }

        Map<String, Byte> map = new HashMap<>();
        for (Map.Entry<Byte, String> entry : huffmanCodes.entrySet()) {
            //将编码表反向
            map.put(entry.getValue(), entry.getKey());
        }
        //创建一个存放byte的集合
        List<Byte> list = new ArrayList<>();
        for (int i = 0; i < stringBuilder.length(); ) {
            int count = 1;
            boolean flag = true;
            Byte b = null;
            while (flag) {
                String key = stringBuilder.substring(i, i + count);
                b = map.get(key);
                if (b == null) {
                    count++;
                } else {
                    flag = false;
                }
            }
            list.add(b);
            i += count; //将j指针移动到count位
        }
        byte[] bytes = new byte[list.size()];
        for (int i = 0; i < bytes.length; i++) {
            bytes[i] = list.get(i);
        }
        return bytes;
    }

    /**
     * 将一个byte转成一个二进制的字符串
     *
     * @param b
     * @param flag true表示是否需要补高位,如果是最后一个字节无需补8位
     * @return byte对应的二进制补码
     * //最后一个字节如果不满8位是正数按原位返回,如果满8位且是负数仍需flag=true
     */
    public static String byteToBitString(boolean flag, byte b) {
        int temp = b;  //将b 转成int类型
        //如果 temp是个正数,补码就是他本身,会出现不够8位的情况
        if (flag || temp == 0) {//如果最后一位恰巧是0的话也要补够8位
            //这个temp可能
            temp |= 256;
            //256 => 0001 0000 0000 => 1 0000 0000
            //1 | 256 => 0000 0001  | 1 0000 0000 =1 0000 0001
        }
        String str = Integer.toBinaryString(temp);  // int的长度有4个字节

        if (flag) {
            return str.substring(str.length() - 8);  //取低8位
        }

        return str;
    }


    //将字节数组转换成Node数组
    public static List<Node> getNodes(byte[] bytes) {
        List<Node> nodes = new ArrayList<>();
        Map<Byte, Integer> map = new HashMap<>();
        for (byte item : bytes) {
            Integer count = map.get(item);
            if (count == null) {
                map.put(item, 1);
            } else {
                map.put(item, count + 1);
            }
        }

        for (Map.Entry<Byte, Integer> entry : map.entrySet()) {
            nodes.add(new Node(entry.getKey(), entry.getValue()));
        }
        return nodes;
    }

    static Map<Byte, String> huffmanCodes = new HashMap<>();
    //static StringBuilder stringBuilder = new StringBuilder();

    /**
     * 将传入的树的所有叶子节点的哈夫曼编码得到,并放入到map集合中
     *
     * @param node
     * @param code
     * @param stringBuilder
     */
    public static void getCodes(Node node, String code, StringBuilder stringBuilder) {
        //获取到上一级的stringBuilder,并创建一个和它有一样值的新StringBuilder
        StringBuilder stringBuilder1 = new StringBuilder(stringBuilder);
        stringBuilder1.append(code);
        if (node != null) {
            if (node.data == null) {
                getCodes(node.left, "0", stringBuilder1);
                getCodes(node.right, "1", stringBuilder1);
            } else {
                huffmanCodes.put(node.data, stringBuilder1.toString());
            }
        }
    }

    //生成编码表方法重载
    public static Map<Byte, String> getCodes(Node node) {
        if (node != null) {
            getCodes(node, "", new StringBuilder());
        } else {
            return null;
        }
        return huffmanCodes;
    }

    /**
     * @param bytes 字符串对应的byte数组
     * @param map   编码表
     * @return 压缩后的byte数组
     * example: content="i like"
     * 返回的是 101 01 000 1111 101 1110
     * 由于计算机底层都是用补码存储的
     * huffmanCodeBytes[0] =10101000(这是补码),我们要传递原码
     * [10101000 => 10101000 - 1 => 10100111(反码)]
     * 10100111  => 11011000  => -88
     * 所以10101000转换成 十六位进制 其实是 -88
     */
    public static byte[] zip(byte[] bytes, Map<Byte, String> map) {
        StringBuilder stringBuilder = new StringBuilder();
        for (byte item : bytes) {
            stringBuilder.append(huffmanCodes.get(item));
        }

        //将stringBuilder按8位一个长度计算
        int len = (stringBuilder.length() + 7) / 8;
        byte[] huffmanCodeBytes = new byte[len];
        int i = 0, index = 0;
        //System.out.println(stringBuilder.substring(0,8)); //10101000
        //将stringBuilder-8位的字符串转换成字节
        if (stringBuilder.length() > 8) {
            for (; i < stringBuilder.length() - 8; i += 8, index++) {
                String strByte = stringBuilder.substring(i, i + 8);
                huffmanCodeBytes[index] = (byte) Integer.parseInt(strByte, 2);
            }
        }
        //将剩余的长度转换成字节
        huffmanCodeBytes[index] = (byte) Integer.parseInt(stringBuilder.substring(i), 2);
        return huffmanCodeBytes;
    }

    //生成哈夫曼树
    public static Node createHuffManTtree(List<Node> nodes) {
        while (nodes.size() > 1) {
            Collections.sort(nodes);
            Node leftNode = nodes.get(0);
            Node rightNode = nodes.get(1);
            Node parent = new Node(null, leftNode.weight + rightNode.weight);
            parent.left = leftNode;
            parent.right = rightNode;
            nodes.remove(leftNode);
            nodes.remove(rightNode);
            nodes.add(parent);
        }
        return nodes.get(0);
    }

    //前序遍历
    public static void preOrder(Node node) {
        if (node != null) {
            System.out.println(node);
            preOrder(node.left);
            preOrder(node.right);
        }
    }
}

class Node implements Comparable<Node> {
    Byte data;    //data代表数据本身
    int weight;   //weight代表权重
    Node left;
    Node right;

    public Node(Byte data, int weight) {
        this.data = data;
        this.weight = weight;
    }

    @Override
    public int compareTo(Node node) {
        return this.weight - node.weight;
    }

    @Override
    public String toString() {
        return "Node{" +
                "data=" + data +
                ", weight=" + weight +
                '}';
    }
}
